package step

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"github.com/wanchain/go-mpc/common/hexutil"
	"github.com/wanchain/go-mpc/log"
	"github.com/wanchain/go-mpc/storeman/osmconf"
	schcomm "github.com/wanchain/go-mpc/storeman/schnorrcomm"
	mpcprotocol "github.com/wanchain/go-mpc/storeman/storemanmpc/protocol"
	"github.com/wanchain/go-mpc/storeman/validator"
	"math/big"
	"strconv"
)

type MpcSSahreJudgeStep struct {
	BaseStep
	SSlshCount uint16
}

func CreateMpcSSahreJudgeStep(totalNum int) *MpcSSahreJudgeStep {
	return &MpcSSahreJudgeStep{
		*CreateBaseStep(totalNum, 0), 0}
}

func (ssj *MpcSSahreJudgeStep) InitStep(result mpcprotocol.MpcResultInterface) error {
	ssj.BaseStep.InitStep(result)
	keyErrNum := mpcprotocol.SShareErrNum
	err := ssj.mpcResult.SetValue(keyErrNum, []big.Int{*big.NewInt(0)})
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep set intializtion of SShareErrNum error", "error", err.Error())
		return err
	}
	return nil
}

func (ssj *MpcSSahreJudgeStep) CreateMessage() []mpcprotocol.StepMessage {

	_, grpIdString, _ := osmconf.GetGrpId(ssj.mpcResult)
	isLeader, _ := osmconf.GetOsmConf().IsLeader(grpIdString)
	log.SyslogDebug("MpcSSahreJudgeStep CreateMessage", "leader ", isLeader, "ctx", ssj.GetMpcContextId())
	if !isLeader {
		// follow no need to wait others judge message. only receive self message
		ssj.SetWaiting(1)
	} else {
		// only check the first sshare err infomation.
		ssj.SetWaiting(1)
	}

	keyErrNum := mpcprotocol.SShareErrNum
	errNum, err := ssj.mpcResult.GetValue(keyErrNum)
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep CreateMessage get SShareErrNum fail", "key", keyErrNum, "ctx", ssj.GetMpcContextId())
	} else {
		log.SyslogDebug("MpcSSahreJudgeStep CreateMessage get SShareErrNum success", "key", keyErrNum, "value", errNum[0].Int64(), "ctx", ssj.GetMpcContextId())
	}
	errNumInt64 := errNum[0].Int64()

	var ret []mpcprotocol.StepMessage
	leaderIndex, _ := osmconf.GetOsmConf().GetLeaderIndex(grpIdString)
	leaderPeerId, _ := osmconf.GetOsmConf().GetNodeIdByIndex(grpIdString, leaderIndex)
	if errNumInt64 > 0 {

		for i := 0; i < int(errNumInt64); i++ {
			ret = make([]mpcprotocol.StepMessage, int(errNumInt64))
			keyErrInfo := mpcprotocol.SShareErrInfos + strconv.Itoa(int(i))
			errInfo, err := ssj.mpcResult.GetValue(keyErrInfo)
			if err != nil {
				log.SyslogErr("MpcSSahreJudgeStep CreateMessage get SShareErrInfos fail", "key", keyErrInfo, "ctx", ssj.GetMpcContextId())
			} else {

				for index, errInfoItem := range errInfo {
					log.SyslogDebug("MpcSSahreJudgeStep CreateMessage get SShareErrInfos success", "key", keyErrInfo,
						"index", index, "info", hexutil.Encode(errInfoItem.Bytes()), "ctx", ssj.GetMpcContextId())
				}

			}

			data := make([]big.Int, 6)
			for j := 0; j < 5; j++ {
				data[j] = errInfo[j]
			}
			// indicate there is  error.
			data[5] = *big.NewInt(128)
			// send multi judge message to leader,since there are more than one error.
			ret[i] = mpcprotocol.StepMessage{MsgCode: mpcprotocol.MPCMessage,
				PeerID:    leaderPeerId,
				Data:      data,
				BytesData: nil}
			// only first error info sent to leader
			break
		}
	} else {
		ret = make([]mpcprotocol.StepMessage, 1)
		data := make([]big.Int, 6)
		for j := 0; j < 5; j++ {
			data[j] = *big.NewInt(1)
		}
		// indicate there is no error.
		data[5] = *big.NewInt(127)
		ret[0] = mpcprotocol.StepMessage{MsgCode: mpcprotocol.MPCMessage,
			PeerID:    leaderPeerId,
			Data:      data,
			BytesData: nil}
		log.SyslogInfo("sshare judge there is NO error record to be judged.", "ctx", ssj.GetMpcContextId())
	}

	if !isLeader {
		// for follow node, send message to self in order to invoke the message to quit itself.
		data := make([]big.Int, 6)
		for j := 0; j < 5; j++ {
			data[j] = *big.NewInt(1)
		}
		// indicate there is no error.
		data[5] = *big.NewInt(127)
		ret = append(ret, mpcprotocol.StepMessage{MsgCode: mpcprotocol.MPCMessage,
			PeerID:    osmconf.GetOsmConf().SelfNodeId,
			Data:      data,
			BytesData: nil})
	}

	return ret
}

func (ssj *MpcSSahreJudgeStep) FinishStep(result mpcprotocol.MpcResultInterface, mpc mpcprotocol.StoremanManager) error {

	err := ssj.BaseStep.FinishStep()
	if err != nil {
		return err
	}

	err = ssj.saveSlshCount(int(ssj.SSlshCount))
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep", "FinishStep err", err.Error(), "ctx", ssj.GetMpcContextId())
		return err
	}

	if ssj.SSlshCount > 0 {
		log.SyslogWarning("MpcSSahreJudgeStep", ":-(:-(:-(FinishStep SSlshCount", ssj.SSlshCount, "ctx", ssj.GetMpcContextId())
		return mpcprotocol.ErrSSlsh
	}
	return nil
}

func (ssj *MpcSSahreJudgeStep) HandleMessage(msg *mpcprotocol.StepMessage) bool {

	if len(msg.Data) != 6 {
		log.SyslogErr("MpcSSahreJudgeStep", "error message recieved. expect len", 6, "Real len", len(msg.Data))
		// sm does not report the error.
		return true
	}
	hasError := int(msg.Data[5].Int64())
	if hasError == 127 {
		//there is no error
		return true
	}

	senderIndex := int(msg.Data[0].Int64())
	rcvIndex := int(msg.Data[1].Int64())
	sshare := msg.Data[2]
	r := msg.Data[3]
	s := msg.Data[4]

	grpId, grpIdString, _ := osmconf.GetGrpId(ssj.mpcResult)

	senderPk, _ := osmconf.GetOsmConf().GetPK(grpIdString, uint16(senderIndex))
	err := schcomm.CheckPK(senderPk)
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep", "HandleMessage", err.Error(), "ctx", ssj.GetMpcContextId())
	}
	// 1. check sig
	h := sha256.Sum256(sshare.Bytes())
	bVerifySig := schcomm.VerifyInternalData(senderPk, h[:], &r, &s)

	if !bVerifySig {
		log.SyslogDebug("MpcSSahreJudgeStep check sig fail")
		log.SyslogErr("MpcSSahreJudgeStep", "senderIndex", senderIndex,
			"rcvIndex", rcvIndex,
			"sshare", hex.EncodeToString(sshare.Bytes()),
			"r", hex.EncodeToString(r.Bytes()),
			"s", hex.EncodeToString(s.Bytes()), "ctx", ssj.GetMpcContextId())

	} else {
		log.SyslogDebug("MpcSSahreJudgeStep check sig successfully", "ctx", ssj.GetMpcContextId())
	}

	// 2. check s content
	rpkShare, _ := ssj.getRPkShare(uint16(senderIndex))
	gpkShare, _ := ssj.getGPKShare(uint16(senderIndex))
	m, _ := ssj.getm()
	bContentCheck, _ := ssj.checkContent(&sshare, m, rpkShare, gpkShare)

	if !bContentCheck {
		log.SyslogDebug("MpcSSahreJudgeStep check content fail", "ctx", ssj.GetMpcContextId())
		log.SyslogErr("MpcSSahreJudgeStep", "senderIndex", senderIndex,
			"rcvIndex", rcvIndex,
			"ctx", ssj.GetMpcContextId())
	} else {
		log.SyslogDebug("MpcSSahreJudgeStep check content successfully", "ctx", ssj.GetMpcContextId())
	}

	log.SyslogDebug("MpcSSahreJudgeStep check result", "checkSig", bVerifySig, "checkContent", bContentCheck, "ctx", ssj.GetMpcContextId())
	bSnderWrong := true

	if !bVerifySig || !bContentCheck {
		bSnderWrong = true
	} else {
		bSnderWrong = false
	}

	// when receive challenge data, there must be one wrong. either send wrong, or receiver wrong.
	// since system can only check the malice by snder, can not make sure malice from reciever, so only slash the sender sm.
	if bSnderWrong {
		ssj.SSlshCount += 1

		ssj.saveSlshProof(bSnderWrong, m, &sshare, &r, &s, senderIndex, rcvIndex, int(ssj.SSlshCount), rpkShare, gpkShare, grpId)
		validator.IncMaliceCount(grpIdString, validator.GetMalicIndex(bSnderWrong, uint16(senderIndex), uint16(rcvIndex)))
	}

	return true
}

func (ssj *MpcSSahreJudgeStep) checkContent(sshare, m *big.Int, rpkShare, gpkShare mpcprotocol.CurvePointer) (bool, error) {

	smpcer := ssj.schnorrMpcer
	sshareG, _ := smpcer.SkG(sshare)
	mPkShare, _ := smpcer.MulPK(m, gpkShare)

	pkTemp, _ := smpcer.Add(rpkShare, mPkShare)

	left := sshareG
	right := pkTemp
	return smpcer.Equal(left, right), nil
}

func (ssj *MpcSSahreJudgeStep) getRPkShare(index uint16) (mpcprotocol.CurvePointer, error) {

	key := mpcprotocol.RPkShare + strconv.Itoa(int(index))
	pkBytes, _ := ssj.mpcResult.GetByteValue(key)
	return ssj.schnorrMpcer.UnMarshPt(pkBytes)
}

func (ssj *MpcSSahreJudgeStep) getm() (*big.Int, error) {
	// check signVerify
	result := ssj.mpcResult
	M, err := result.GetByteValue(mpcprotocol.MpcM)
	if err != nil {
		log.SyslogErr("MpcAckRSStep::verifyRS", "ack MpcAckRSStep get MpcM . err", err.Error(), "ctx", ssj.GetMpcContextId())
		return &big.Int{}, err
	}

	hashMBytes := sha256.Sum256(M)

	// rpk : R
	rpkBytes, _ := result.GetByteValue(mpcprotocol.RPk)

	// Forming the m: hash(message||rpk)
	var buffer bytes.Buffer
	//buffer.Write(M[:])
	buffer.Write(hashMBytes[:])
	buffer.Write(rpkBytes)
	mTemp := sha256.Sum256(buffer.Bytes())
	m := new(big.Int).SetBytes(mTemp[:])
	m = m.Mod(m, ssj.schnorrMpcer.GetMod())
	return m, nil
}

func (ssj *MpcSSahreJudgeStep) getGPKShare(index uint16) (mpcprotocol.CurvePointer, error) {
	//

	_, grpIdString, err := osmconf.GetGrpId(ssj.mpcResult)
	if err != nil {
		return nil, err
	}

	_, gpkStr, err := osmconf.GetGrpGpk(ssj.mpcResult)
	if err != nil {
		return nil, err
	}

	gpkShareBytes, err := osmconf.GetOsmConf().GetPKShareBytes(grpIdString, gpkStr, index, ssj.CurveType())
	if err != nil {
		return nil, err
	}
	return ssj.schnorrMpcer.UnMarshPt(gpkShareBytes)
}

func (ssj *MpcSSahreJudgeStep) saveSlshCount(slshCount int) error {

	sslshValue := make([]big.Int, 1)
	sslshValue[0] = *big.NewInt(0).SetInt64(int64(ssj.SSlshCount))

	key := mpcprotocol.MPCSSlshProofNum
	err := ssj.mpcResult.SetValue(key, sslshValue)
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep", "save MPCSSlshProofNum", err.Error(), "key", key, "ctx", ssj.GetMpcContextId())
		return err
	} else {
		if slshCount > 0 {
			log.SyslogWarning("MpcSSahreJudgeStep", "save MPCSSlshProofNum success key", key, "slshCount", slshCount, "ctx", ssj.GetMpcContextId())
		}
	}

	return nil
}

func (ssj *MpcSSahreJudgeStep) saveSlshProof(isSnder bool,
	m, sshare, r, s *big.Int,
	sndrIndex, rcvrIndex, slshCount int,
	rpkShare, gpkShare mpcprotocol.CurvePointer, grp []byte) error {

	sslshValue := make([]big.Int, 7)
	if isSnder {
		sslshValue[0] = *schcomm.BigOne
	} else {
		sslshValue[0] = *schcomm.BigZero
	}

	sslshValue[1] = *m
	sslshValue[2] = *sshare
	sslshValue[3] = *r
	sslshValue[4] = *s
	sslshValue[5] = *big.NewInt(0).SetInt64(int64(sndrIndex))
	sslshValue[6] = *big.NewInt(0).SetInt64(int64(rcvrIndex))

	smpcer := ssj.schnorrMpcer
	// rpkShare, gpkShare, grpId
	var sslshByte bytes.Buffer
	rpkShareBytes, err := smpcer.MarshPt(rpkShare)
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep", "MarshPt(rpkShare) err ", err.Error(), "ctx", ssj.GetMpcContextId())
		return err
	}
	gpkShareBytes, err := smpcer.MarshPt(gpkShare)
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep", "MarshPt(gpkShare) err ", err.Error(), "ctx", ssj.GetMpcContextId())
		return err
	}
	sslshByte.Write(rpkShareBytes)
	sslshByte.Write(gpkShareBytes)
	sslshByte.Write(grp[:])

	key1 := mpcprotocol.SSlshProof + strconv.Itoa(int(slshCount-1))
	err = ssj.mpcResult.SetValue(key1, sslshValue)
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep", "save SlshProof.SetValue err ", err.Error(), "key", key1, "ctx", ssj.GetMpcContextId())
		return err
	} else {
		log.SyslogErr("MpcSSahreJudgeStep", "save SlshProof.SetValue success ", "key", key1, "ctx", ssj.GetMpcContextId())
	}
	err = ssj.mpcResult.SetByteValue(key1, sslshByte.Bytes())
	if err != nil {
		log.SyslogErr("MpcSSahreJudgeStep", "saveSlshProof.SetByteValue err ", err.Error(), "key", key1, "ctx", ssj.GetMpcContextId())
		return err
	} else {
		log.SyslogErr("MpcSSahreJudgeStep", "save SlshProof.SetByteValue success ", "key", key1, "ctx", ssj.GetMpcContextId())
	}

	return nil
}
